import numpy as np
from fig_utils import Util, result_fig
from load_data import dataset


def _euclidean_distance(v1, v2):
    """
    Compute Euclidean distance between v1 and v2
    :param v1: a point vector
    :param v2: a point vector
    :return:
    """
    return np.sqrt(sum((v2 - v1) ** 2))


class Kmeans(object):

    def __init__(self, data, num_class):
        self._data = data
        self._num_class = num_class
        self._N, self._d = data.shape
        self._centroids = np.zeros((num_class, self._d))
        self._w = []

    def init_centroids(self):
        """
        Init centroids
        :return:
        """
        for i in range(self._num_class):
            index = np.random.randint(0, self._N)
            self._centroids[i, :] = self._data[index, :]

        return

    def _wcss(self, clusters):
        """
        Compute WCSS (within cluster sum of squares):
        :param clusters:
        :return:
        """
        result = 0
        for i in clusters[:, 1]:
            result += i ** 2
        self._w.append(result)

    def fit(self):
        """
        Perform k-means
        :return:
        """
        cluster = np.zeros((self._N, 2))
        changed = True
        self.init_centroids()
        iter = 0
        while changed:
            changed = False

            min_index = 0

            for i in range(self._N):
                min_dist = np.inf
                for j in range(self._num_class):
                    distance = _euclidean_distance(self._centroids[j, :], self._data[i, :])
                    if distance < min_dist:
                        min_dist = distance
                        cluster[i, 1] = min_dist
                        min_index = j

                if cluster[i, 0] != min_index:
                    changed = True
                    cluster[i, 0] = min_index

            predicted_label = cluster[:, 0].astype(np.int)

            for j in range(self._num_class):
                cluster_index = np.nonzero(cluster[:, 0] == j)
                points_in_cluster = self._data[cluster_index]
                self._centroids[j, :] = np.mean(points_in_cluster, axis=0)

            self._wcss(cluster)

            result_fig(self._data, self._w, iter, 100, self._num_class,
                       label=predicted_label, mean=self._centroids, interval=0.3, kmeans=True, end=not changed)
            iter += 1

        return


if __name__ == '__main__':
    num_class = 6
    _, data = dataset.load_data(num_class)
    model = Kmeans(data, num_class)
    model.fit()
